import sys
import fitz  # PyMuPDF
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QPushButton, QFileDialog, QInputDialog,
                             QLabel, QScrollArea, QMessageBox)
from PyQt6.QtGui import QPixmap, QImage, QMouseEvent
from PyQt6.QtCore import Qt, QPoint


class ClickableLabel(QLabel):
    """支持鼠标点击事件的标签控件"""
    clicked = Qt.pyqtSignal(QPoint)

    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.MouseButton.LeftButton:
            self.clicked.emit(event.pos())
        super().mousePressEvent(event)


class PDFEditor(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("PDF编辑器")
        self.setGeometry(100, 100, 1000, 700)

        self.pdf_document = None
        self.pages = []  # 存储(QLabel, fitz.Page, 缩放因子)
        self.annotations = []  # 存储(page_num, type, content, rect)
        self.current_editing_page = -1  # 当前编辑的页码

        self.init_ui()

    def init_ui(self):
        # 菜单栏
        menubar = self.menuBar()
        file_menu = menubar.addMenu("文件")
        open_action = file_menu.addAction("打开PDF")
        open_action.triggered.connect(self.open_pdf)
        save_action = file_menu.addAction("保存PDF")
        save_action.triggered.connect(self.save_pdf)

        # 工具栏
        tool_layout = QHBoxLayout()
        self.add_text_btn = QPushButton("添加文本")
        self.add_text_btn.clicked.connect(self.enable_text_editing)
        tool_layout.addWidget(self.add_text_btn)

        # 状态标签
        self.status_label = QLabel("就绪")
        tool_layout.addWidget(self.status_label)
        tool_layout.addStretch()

        # PDF预览区域
        self.scroll_area = QScrollArea()
        self.scroll_content = QWidget()
        self.scroll_layout = QVBoxLayout(self.scroll_content)
        self.scroll_area.setWidget(self.scroll_content)
        self.scroll_area.setWidgetResizable(True)

        # 主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        main_layout.addLayout(tool_layout)
        main_layout.addWidget(self.scroll_area)

    def open_pdf(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "打开PDF", "", "PDF Files (*.pdf)")
        if file_path:
            try:
                self.pdf_document = fitz.open(file_path)
                self._render_pages()
                QMessageBox.information(self, "成功", f"已打开PDF文件，共{len(self.pdf_document)}页")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"打开PDF失败: {str(e)}")

    def _render_pages(self):
        # 清除现有内容
        while self.scroll_layout.count():
            item = self.scroll_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()

        self.pages = []
        for page_num in range(len(self.pdf_document)):
            page = self.pdf_document.load_page(page_num)
            pix = page.get_pixmap(matrix=fitz.Matrix(1.5, 1.5))  # 放大1.5倍以便编辑

            # 修复QImage格式问题
            try:
                img = QImage(pix.samples, pix.width, pix.height, pix.stride, QImage.Format.Format_RGB888)
            except AttributeError:
                img = QImage(pix.samples, pix.width, pix.height, pix.stride, QImage.Format.Format_RGB32)

            # 使用自定义的可点击标签
            label = ClickableLabel()
            label.setPixmap(QPixmap.fromImage(img))
            label.clicked.connect(lambda pos, pn=page_num: self.on_page_clicked(pos, pn))
            self.scroll_layout.addWidget(label)

            # 存储页面信息和缩放因子
            self.pages.append((label, page, 1.5))

    def enable_text_editing(self):
        if not self.pages:
            QMessageBox.warning(self, "警告", "请先打开PDF文件")
            return

        self.current_editing_page = -1  # 重置当前编辑页
        self.status_label.setText("请点击PDF页面选择文本添加位置...")
        QMessageBox.information(self, "提示", "请在PDF页面上点击要添加文本的位置")

    def on_page_clicked(self, pos, page_num):
        # 如果不是处于编辑模式，不处理点击
        if self.status_label.text() != "请点击PDF页面选择文本添加位置...":
            return

        self.current_editing_page = page_num
        self.status_label.setText(f"已选择页面 {page_num + 1} 的位置，请输入文本")

        # 获取点击位置（需要根据缩放因子转换）
        label, page, zoom = self.pages[page_num]
        x = pos.x() / zoom
        y = pos.y() / zoom

        # 请求用户输入文本
        text, ok = QInputDialog.getText(self, "添加文本", "输入文本内容:")
        if ok and text:
            try:
                # 在点击位置添加文本
                rect = fitz.Rect(x, y, x + 150, y + 20)  # 宽度150，高度20的文本框
                self.annotations.append((page_num, "text", text, rect))
                self._update_preview()
                QMessageBox.information(self, "成功", "文本已添加，请保存PDF查看效果")
                self.status_label.setText("就绪")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"添加文本失败: {str(e)}")
                self.status_label.setText("就绪")

    def _update_preview(self):
        try:
            for page_num in range(len(self.pages)):
                label, page, zoom = self.pages[page_num]
                temp_page = page.copy()

                # 应用该页的所有注释
                for annot in self.annotations:
                    a_page_num, a_type, content, rect = annot
                    if a_page_num == page_num and a_type == "text":
                        # 添加红色文本
                        temp_page.insert_text(rect.tl, content, fontsize=12, color=(1, 0, 0))

                # 重新渲染页面
                pix = temp_page.get_pixmap(matrix=fitz.Matrix(zoom, zoom))
                try:
                    img = QImage(pix.samples, pix.width, pix.height, pix.stride, QImage.Format.Format_RGB888)
                except AttributeError:
                    img = QImage(pix.samples, pix.width, pix.height, pix.stride, QImage.Format.Format_RGB32)

                label.setPixmap(QPixmap.fromImage(img))
        except Exception as e:
            QMessageBox.critical(self, "错误", f"更新预览失败: {str(e)}")

    def save_pdf(self):
        if not self.pdf_document:
            QMessageBox.warning(self, "警告", "请先打开PDF文件")
            return

        file_path, _ = QFileDialog.getSaveFileName(self, "保存PDF", "", "PDF Files (*.pdf)")
        if file_path:
            try:
                new_pdf = fitz.open()
                for page_num in range(len(self.pdf_document)):
                    page = self.pdf_document.load_page(page_num)
                    # 应用注释
                    for annot in self.annotations:
                        a_page_num, a_type, content, rect = annot
                        if a_page_num == page_num and a_type == "text":
                            page.insert_text(rect.tl, content, fontsize=12)
                    new_pdf.insert_pdf(self.pdf_document, from_page=page_num, to_page=page_num)

                new_pdf.save(file_path)
                new_pdf.close()
                QMessageBox.information(self, "成功", f"PDF文件已保存至:\n{file_path}")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存PDF失败: {str(e)}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = PDFEditor()
    window.show()
    sys.exit(app.exec())